package com.study.graph.config;

import com.vesoft.nebula.client.graph.NebulaPoolConfig;
import com.vesoft.nebula.client.graph.SessionPool;
import com.vesoft.nebula.client.graph.SessionPoolConfig;
import com.vesoft.nebula.client.graph.data.HostAddress;
import com.vesoft.nebula.client.graph.net.NebulaPool;
import com.vesoft.nebula.client.graph.net.Session;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Scope;
import org.springframework.context.annotation.ScopedProxyMode;

import java.util.List;
import java.util.stream.Collectors;

/**
 * NebulaGraph Java 客户端提供 Connection Pool 和 Session Pool 两种使用方式
 */
@Slf4j
@Configuration
public class NebulaConfig {

    @Bean
    public NebulaPool nebulaPool(NebulaProperties nebulaProperties) throws Exception {
        NebulaPool pool = new NebulaPool();
        NebulaPoolConfig nebulaPoolConfig = new NebulaPoolConfig();
        nebulaPoolConfig.setMaxConnSize(1000);
        List<HostAddress> addressList = nebulaProperties.getAddress().stream()
                .map(d -> new HostAddress(d.getHost(), d.getPort()))
                .collect(Collectors.toList());
        boolean init = pool.init(addressList, nebulaPoolConfig);
        if (init) {
            log.info("NebulaGraph init Success ！");
            return pool;
        }
        throw new RuntimeException("NebulaGraph init err !");
    }

    @Bean
    @Scope(scopeName = "prototype", proxyMode = ScopedProxyMode.TARGET_CLASS)
    public Session session(NebulaPool nebulaPool, NebulaProperties properties) {
        try {
            Session session = nebulaPool.getSession(properties.getUsername(), properties.getPassword(), properties.isReconnect());
            session.execute(NebulaConstant.USE.concat(properties.getSpace()).concat(NebulaConstant.SEMICOLON));
            return session;
        } catch (Exception e) {
            log.error("get nebula session err , {} ", e.toString());
            return null;
        }
    }

    @Bean
    public SessionPool sessionPool(NebulaProperties properties){
        List<HostAddress> address = properties.getAddress().stream()
                .map(info-> new HostAddress(info.getHost(),info.getPort()))
                .collect(Collectors.toList());
        SessionPoolConfig sessionPoolConfig =
                new SessionPoolConfig(address, properties.getSpace(), properties.getUsername(), properties.getPassword())
                        .setMaxSessionSize(10)
                        .setMinSessionSize(10)
                        .setRetryConnectTimes(3)
                        .setWaitTime(100)
                        .setRetryTimes(3)
                        .setIntervalTime(100);
        SessionPool sessionPool = new SessionPool(sessionPoolConfig);
        if (sessionPool.init()) {
            return sessionPool;
        }
        log.error("session pool init failed.");
        throw new RuntimeException("NebulaGraph init err !");
    }
}
